/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.runners.gaia.translation.functions;



import com.bff.gaia.unified.runners.core.ReduceFnRunner;

import com.bff.gaia.unified.runners.core.SideInputReader;

import com.bff.gaia.unified.sdk.options.PipelineOptions;

import com.bff.gaia.unified.sdk.transforms.windowing.BoundedWindow;

import com.bff.gaia.unified.sdk.transforms.windowing.PaneInfo;

import com.bff.gaia.unified.sdk.transforms.windowing.TimestampCombiner;

import com.bff.gaia.unified.sdk.transforms.windowing.WindowFn;

import com.bff.gaia.unified.sdk.util.WindowedValue;

import com.bff.gaia.unified.sdk.values.KV;

import com.bff.gaia.unified.sdk.values.WindowingStrategy;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.Iterables;

import com.bff.gaia.api.java.tuple.Tuple2;

import com.bff.gaia.util.Collector;

import org.joda.time.Instant;



import java.util.*;



/**

 * A Gaia combine runner that builds a map of merged windows and produces output after seeing all

 * input. This is similar to what{@link ReduceFnRunner} does.

 */

public class HashingGaiaCombineRunner<K, InputT, AccumT, OutputT, W extends BoundedWindow>

    extends AbstractGaiaCombineRunner<K, InputT, AccumT, OutputT, W> {



  @Override

  public void combine(

      GaiaCombiner<K, InputT, AccumT, OutputT> gaiaCombiner,

      WindowingStrategy<Object, W> windowingStrategy,

      SideInputReader sideInputReader,

      PipelineOptions options,

      Iterable<WindowedValue<KV<K, InputT>>> elements,

      Collector<WindowedValue<KV<K, OutputT>>> out)

      throws Exception {



    @SuppressWarnings("unchecked")

    TimestampCombiner timestampCombiner = windowingStrategy.getTimestampCombiner();

    WindowFn<Object, W> windowFn = windowingStrategy.getWindowFn();



    // Gaia Iterable can be iterated over only once.

    List<WindowedValue<KV<K, InputT>>> inputs = new ArrayList<>();

    Iterables.addAll(inputs, elements);



    Set<W> windows = collectWindows(inputs);

    Map<W, W> windowToMergeResult = mergeWindows(windowingStrategy, windows);



    // Combine all windowedValues into map

    Map<W, Tuple2<AccumT, Instant>> mapState = new HashMap<>();

    Iterator<WindowedValue<KV<K, InputT>>> iterator = inputs.iterator();

    WindowedValue<KV<K, InputT>> currentValue = iterator.next();

    K key = currentValue.getValue().getKey();

    do {

      for (BoundedWindow w : currentValue.getWindows()) {

        @SuppressWarnings("unchecked")

        W currentWindow = (W) w;

        W mergedWindow = windowToMergeResult.get(currentWindow);

        mergedWindow = mergedWindow == null ? currentWindow : mergedWindow;

        Set<W> singletonW = Collections.singleton(mergedWindow);

        Tuple2<AccumT, Instant> accumAndInstant = mapState.get(mergedWindow);

        if (accumAndInstant == null) {

          AccumT accumT =

              gaiaCombiner.firstInput(

                  key, currentValue.getValue().getValue(), options, sideInputReader, singletonW);

          Instant windowTimestamp =

              timestampCombiner.assign(

                  mergedWindow, windowFn.getOutputTime(currentValue.getTimestamp(), mergedWindow));

          accumAndInstant = new Tuple2<>(accumT, windowTimestamp);

          mapState.put(mergedWindow, accumAndInstant);

        } else {

          accumAndInstant.f0 =

              gaiaCombiner.addInput(

                  key,

                  accumAndInstant.f0,

                  currentValue.getValue().getValue(),

                  options,

                  sideInputReader,

                  singletonW);

          accumAndInstant.f1 =

              timestampCombiner.combine(

                  accumAndInstant.f1,

                  timestampCombiner.assign(

                      mergedWindow,

                      windowingStrategy

                          .getWindowFn()

                          .getOutputTime(currentValue.getTimestamp(), mergedWindow)));

        }

      }

      if (iterator.hasNext()) {

        currentValue = iterator.next();

      } else {

        break;

      }

    } while (true);



    // Output the final value of combiners

    for (Map.Entry<W, Tuple2<AccumT, Instant>> entry : mapState.entrySet()) {

      AccumT accumulator = entry.getValue().f0;

      Instant windowTimestamp = entry.getValue().f1;

      out.collect(

          WindowedValue.of(

              KV.of(

                  key,

                  gaiaCombiner.extractOutput(

                      key,

                      accumulator,

                      options,

                      sideInputReader,

                      Collections.singleton(entry.getKey()))),

              windowTimestamp,

              entry.getKey(),

              PaneInfo.NO_FIRING));

    }

  }



  private Map<W, W> mergeWindows(WindowingStrategy<Object, W> windowingStrategy, Set<W> windows)

      throws Exception {

    WindowFn<Object, W> windowFn = windowingStrategy.getWindowFn();



    if (windowingStrategy.getWindowFn().isNonMerging()) {

      // Return an empty map, indicating that every window is not merged.

      return Collections.emptyMap();

    }



    Map<W, W> windowToMergeResult = new HashMap<>();

    windowFn.mergeWindows(new MergeContextImpl(windowFn, windows, windowToMergeResult));

    return windowToMergeResult;

  }



  private class MergeContextImpl extends WindowFn<Object, W>.MergeContext {



    private Set<W> windows;

    private Map<W, W> windowToMergeResult;



    MergeContextImpl(WindowFn<Object, W> windowFn, Set<W> windows, Map<W, W> windowToMergeResult) {

      windowFn.super();

      this.windows = windows;

      this.windowToMergeResult = windowToMergeResult;

    }



    @Override

    public Collection<W> windows() {

      return windows;

    }



    @Override

    public void merge(Collection<W> toBeMerged, W mergeResult) throws Exception {

      for (W w : toBeMerged) {

        windowToMergeResult.put(w, mergeResult);

      }

    }

  }



  private Set<W> collectWindows(Iterable<WindowedValue<KV<K, InputT>>> values) {

    Set<W> windows = new HashSet<>();

    for (WindowedValue<?> value : values) {

      for (BoundedWindow untypedWindow : value.getWindows()) {

        @SuppressWarnings("unchecked")

        W window = (W) untypedWindow;

        windows.add(window);

      }

    }

    return windows;

  }

}